package org.jsoncrypto;

import java.security.SecureRandom;

import javax.crypto.SecretKey;

import org.json.JSONException;
import org.json.JSONObject;
import org.jsoncrypto.util.Base64;

abstract class JcLibEnc {
  JSONObject mHeader;
  String mHeaderStr = null;

  SecretKey mContentEncryptionKey = null;
  int mKeyLength = 0;
  String mEncStr = null;

  String mEncodedJwtHeaderSegment = null;
  String mEncodedJwtKeySegment = null;
  String mEncodedJwtIvSegment = null;
  String mEncodedJwtCryptoSegment = null;
  String mEncodedJwtIntegritySegment = null;

  SecureRandom secureRandom = null;

  /*
   * value is only set after a call to encrypt
   */
  String getEncodedJwtCryptoSegment() {
    return mEncodedJwtCryptoSegment;
  }

  /*
   * value is only set after a call to encrypt
   */
  public String getEncodedIntegritySegment() {
    return mEncodedJwtIntegritySegment;
  }

  void setIv(byte[] ivBytes) throws JcBaseException {
    String iv = mHeader.optString("iv");
    if (iv != null) {
      throw new JcBaseException("invalid try to set iv");
    }
    try {
      mHeader.put("iv",
          Base64.encodeBytes(ivBytes, org.jsoncrypto.util.Base64.DONT_BREAK_LINES | org.jsoncrypto.util.Base64.URL));
      mHeaderStr = mHeader.toString();
    } catch (JSONException e) {
      throw new JcBaseException(e);
    }
  }

  void generateRandomIv() {
    // FIXME
  }

  protected JcLibEnc(JSONObject header, String encStr) throws JcBaseException {
    try {
      mEncStr = encStr;
      mHeader = new JSONObject(header.toString());
      String jweEncStr = mHeader.optString("enc", null);
      if (jweEncStr == null) {
        mHeader.put("enc", encStr);
      } else {
        if (!jweEncStr.equals(encStr)) {
          throw new JcBaseException("expected enc " + encStr);
        }
      }
      mHeaderStr = mHeader.toString();
      mKeyLength = JcBase.getEncKeyLength(mEncStr);
    } catch (JSONException e) {
      throw new JcBaseException(e);
    }
  }

  public JcLibEnc(String encodedJwtHeaderSegment, String encodedJwtKeySegment, String encodedJwtIvSegment,
      String encodedJwtCryptoSegment, String encodedJwtIntegritySegment) throws JcBaseException {
    try {
      mEncodedJwtHeaderSegment = encodedJwtHeaderSegment;
      mEncodedJwtKeySegment = encodedJwtKeySegment;
      mEncodedJwtIvSegment = encodedJwtIvSegment;
      mEncodedJwtCryptoSegment = encodedJwtCryptoSegment;
      mEncodedJwtIntegritySegment = encodedJwtIntegritySegment;

      mHeaderStr = new String(Base64.decodeUrl(encodedJwtHeaderSegment));
      mHeader = new JSONObject(mHeaderStr);
      mEncStr = mHeader.getString("enc");
      mKeyLength = JcBase.getEncKeyLength(getEncStr());
    } catch (Exception e) {
      throw new JcBaseException(e);
    }
  }

  String getEncStr() throws JcBaseException {
    return mEncStr;
  }

  int getKeyLength() throws JcBaseException {
    return mKeyLength;
  }

  void setContentEncryptionKey(SecretKey contentEncryptionKey) throws JcBaseException {
    mContentEncryptionKey = contentEncryptionKey;
  }

  /*
   * Generate a symmetric key with a key length suitable for this encryption
   * algorithm
   */
  SecretKey generateRandomContentEncryptionKey() throws JcBaseException {
    if (secureRandom == null) {
      secureRandom = new SecureRandom();
    }
    mContentEncryptionKey = JcBase.generateRandomContentEncryptionKey(getKeyLength(), secureRandom);
    return mContentEncryptionKey;
  }

  String getHeaderStr() {
    return mHeaderStr;
  }

  String getEncodedJweEncryptedKey() {
    return mEncodedJwtKeySegment;
  }

  void setEncodedJweEncryptedKey(String encodedJwtKeySegment) {
    mEncodedJwtKeySegment = encodedJwtKeySegment;
  }

  /*
   * uses the content encryption key to encrypt the content bytes
   */
  abstract String encrypt(byte[] contentBytes, SecretKey contentEncryptionKey, String encodedJweEncryptedKey,
      String encodedJweIvSegment) throws JcBaseException;

  byte[] decrypt(byte[] secretKeyBytes) throws JcBaseException {
    try {
      byte[] jwtIntegritySegmentBytes = Base64.decodeUrl(mEncodedJwtIntegritySegment);
      return JcBase.doDecrypt(mEncodedJwtHeaderSegment, mEncodedJwtKeySegment, mEncodedJwtCryptoSegment,
          jwtIntegritySegmentBytes, mHeader, secretKeyBytes, mEncodedJwtIvSegment);
    } catch (Exception e) {
      throw new JcBaseException(e);
    }
  }

  public static JcLibEnc getInstance(String jwe) throws JcBaseException {
    try {
      String[] split = jwe.split("\\.");
      String encodedJwtHeaderSegment = split[0];
      String encodedJwtKeySegment = split[1];
      String encodedJwtIvSegment = split[2];
      String encodedJwtCryptoSegment = split[3];
      String encodedJwtIntegritySegment = split[4];

      String headerStr = new String(Base64.decodeUrl(encodedJwtHeaderSegment));
      JSONObject header = new JSONObject(headerStr);
      String encStr = header.getString("enc");
      if (JcBase.ENC_ALG_A128GCM.equals(encStr)) {
        return new JcEnc_A128GCM(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A192GCM.equals(encStr)) {
        return new JcEnc_A192GCM(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A256GCM.equals(encStr)) {
        return new JcEnc_A256GCM(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A512GCM.equals(encStr)) {
        return new JcEnc_A512GCM(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A128CBCHS256.equals(encStr)) {
        return new JcEnc_A128CBCHS256(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A192CBCHS256.equals(encStr)) {
        return new JcEnc_A128CBCHS256(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A256CBCHS256.equals(encStr)) {
        return new JcEnc_A128CBCHS256(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      if (JcBase.ENC_ALG_A512CBCHS256.equals(encStr)) {
        return new JcEnc_A512CBCHS256(encodedJwtHeaderSegment, encodedJwtKeySegment, encodedJwtIvSegment,
            encodedJwtCryptoSegment, encodedJwtIntegritySegment);
      }
      throw new JcBaseException("unsupported enc algorithm: " + encStr);
    } catch (Exception e) {
      throw new JcBaseException(e);
    }
  }

  void setAlgStr(String alg) throws JcBaseException {
    if (mHeader.optString("alg", null) != null) {
      throw new JcBaseException("alg is already set");
    }
    try {
      mHeader.put("alg", alg);
    } catch (JSONException e) {
      throw new JcBaseException(e);
    }
    mHeaderStr = mHeader.toString();
  }

  void setEncStr(String alg) throws JcBaseException {
    if (mHeader.optString("enc", null) != null) {
      throw new JcBaseException("enc is already set");
    }
    try {
      mHeader.put("enc", alg);
    } catch (JSONException e) {
      throw new JcBaseException(e);
    }
    mHeaderStr = mHeader.toString();
  }

}
